package com.rankeiot.ai.util;

import cn.hutool.http.HttpUtil;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.rankeiot.ai.domain.Application;
import com.rankeiot.ai.message.ChatCompletionRequest;
import com.rankeiot.ai.message.ChatMessage;
import com.rankeiot.ai.message.FunctionCall;
import com.rankeiot.core.util.JSONUtil;
import com.rankeiot.core.util.StringUtil;
import org.checkerframework.checker.units.qual.C;


import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

public class ErnieClient implements AiClient {
    private String baseUrl;
    private String aiKey;
    private String model;
    private Application application;
    private static Cache<String, String> cache = Caffeine.newBuilder().expireAfterWrite(15, TimeUnit.DAYS).build();

    public ErnieClient(Application application) {
        this.baseUrl = application.getAiHost();
        this.aiKey = application.getAiKey();
        this.model = application.getMode();
        this.application = application;
    }

    public String chatCompletion(ChatCompletionRequest request, BiFunction<String,Map<String, Object>, FunctionResult> funCall) {
        String accessToken = cache.get(aiKey, this::getAccessToken);
        String url = "";
        if ("ERNIE-Bot".equals(model)) {
            url = baseUrl + "rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=";
        } else if ("ERNIE-Bot 4.0".equals(model)) {
            url = baseUrl + "rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=";
        }

        copyIfNotNull(application::getPrompt,request::setSystem);
        copyIfNotNull(application::getTemperature,request::setTemperature);
        copyIfNotNull(application::getTopP,request::setTop_p);
        copyIfNotNull(application::getPenaltyScore,request::setPenalty_score);
        copyIfNotNull(application::getDisableSearch,request::setDisable_search);


        while (true) {
            String resp = HttpUtil.post(url + accessToken, JSONUtil.toJSONString(request));
            System.out.println(resp);
            Map answer = JSONUtil.parseObject(resp);
            if ("normal".equals(answer.get("finish_reason"))) {
                return (String) answer.get("result");
            } else if ("function_call".equals(answer.get("finish_reason"))) {
                Map call = (Map) answer.get("function_call");
                String name = (String) call.get("name");
                String args = (String) call.get("arguments");
                Map argMap = new HashMap();
                if (StringUtil.isNotEmpty(args)) {
                    argMap = JSONUtil.parseObject(args);
                }
                FunctionResult result = funCall.apply(name, argMap);
                if (result == null) {
                    return "function not found";
                }
                if (result.getDirect()) {
                    return result.getData();
                } else {
                    ChatMessage assistantMessage = new ChatMessage();
                    assistantMessage.setRole("assistant");
                    assistantMessage.setContent(null);
                    assistantMessage.setFunction_call(new FunctionCall(name, args));
                    request.getMessages().add(assistantMessage);

                    ChatMessage functionMessage = new ChatMessage();
                    functionMessage.setRole("function");
                    functionMessage.setName(name);
                    functionMessage.setContent(result.getData());
                    request.getMessages().add(functionMessage);
                }
            }
        }
    }
    private static <T> void copyIfNotNull(Supplier<T> supplier, Consumer<T> setter) {
       T object= supplier.get();
       if(object != null){
           setter.accept(object);
       }
    }
    private String getAccessToken(String aiKey) {
        String param = aiKey.replaceFirst("-", "&client_secret=");
        String resp = HttpUtil.get(baseUrl + "oauth/2.0/token?grant_type=client_credentials&client_id=" + param);
        Map<String, String> ret = JSONUtil.parseObject(resp);
        return ret.get("access_token");
    }
}
